package com.gitee.wsl.flow.transform.map

import kotlinx.atomicfu.atomic
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineStart
import kotlinx.coroutines.async
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.FlowCollector
import kotlinx.coroutines.flow.buffer
import kotlinx.coroutines.flow.catch
import kotlinx.coroutines.flow.channelFlow
import kotlinx.coroutines.flow.collectIndexed
import kotlinx.coroutines.flow.emitAll
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.launch

/**
 * Returns a flow containing the results of applying the given [transform] function
 * to each value and its index in the original flow.
 */
fun <T, R> Flow<T>.mapIndexed(transform: suspend (index: Int, value: T) -> R): Flow<R> =
    flow {
        collectIndexed { index, value ->
            return@collectIndexed emit(transform(index, value))
        }
    }

fun <T, R> Flow<T>.map(transform: suspend ( value: T) -> R): Flow<R> =
    flow {
        collect{value ->
            return@collect emit(transform(value))
        }
    }


/**
 * Projects each source value to a [Flow] which is merged in the output [Flow] only if the previous projected [Flow] has completed.
 * If value is received while there is some projected [Flow] sequence being merged, it will simply be ignored.
 *
 * This method is a shortcut for `map(transform).flattenFirst()`. See [flattenFirst].
 *
 * ### Operator fusion
 *
 * Applications of [flowOn], [buffer], and [produceIn] _after_ this operator are fused with
 * its concurrent merging so that only one properly configured channel is used for execution of merging logic.
 *
 * @param transform A transform function to apply to value that was observed while no Flow is executing in parallel.
 */
fun <T, R> Flow<T>.flatMapFirst(transform: suspend (value: T) -> Flow<R>): Flow<R> =
    map(transform).flattenFirst()

/**
 * Converts a higher-order [Flow] into a first-order [Flow] by dropping inner [Flow] while the previous inner [Flow] has not yet completed.
 */
fun <T> Flow<Flow<T>>.flattenFirst(): Flow<T> = channelFlow {
    val busy = atomic(false)

    collect { inner ->
        if (busy.compareAndSet(expect = false, update = true)) {
            // Do not pay for dispatch here, it's never necessary
            launch(start = CoroutineStart.UNDISPATCHED) {
                try {
                    inner.collect { send(it) }
                    busy.value = false
                } catch (e: CancellationException) {
                    busy.value = false
                }
            }
        }
    }
}

/**
 * This function is an alias to [flattenFirst] operator.
 */
@Suppress("NOTHING_TO_INLINE")
inline fun <T> Flow<Flow<T>>.exhaustAll(): Flow<T> = flattenFirst()

/**
 * This function is an alias to [flatMapFirst] operator.
 */
fun <T, R> Flow<T>.exhaustMap(transform: suspend (value: T) -> Flow<R>): Flow<R> =
    flatMapFirst(transform)

fun <T> Flow<T>.catchMap(transform: FlowCollector<T>.(Throwable) -> T): Flow<T> =
    catch { emit(transform(it)) }

fun <T, R> Flow<T>.concurrentMap(
    capacity: Int = Channel.BUFFERED,
    transformBlock: suspend (T) -> R
): Flow<R> {
    return flow {
        coroutineScope {
            map {
                async(start = CoroutineStart.LAZY) {
                    transformBlock(
                        it
                    )
                }
            }
                .buffer(capacity)
                .map {
                    it.start()
                    it.await()
                }
                .let {
                    emitAll(it)
                }
        }
    }
}


/**
 * Maps values of each collection [this] flow emits.
 */
inline fun <T, R> Flow<Iterable<T>>.mapValues(
    crossinline transform: (T) -> R
): Flow<List<R>> = map { it.map(transform) }
